{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Preference Optimization (DPO) at scale with QLoRA\n",
    "\n",
    "This guide provides a step-by-step workflow for preference fine-tuning the [`Qwen/Qwen2.5-7B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) model on a multi-GPU Anyscale cluster. You use LLaMA-Factory as the training framework and `QLoRA` to reduce memory requirements and enable efficient multi-GPU training.\n",
    "\n",
    "DPO aligns a model with human preferences using pairs of “chosen” and “rejected” responses. Rather than training a separate reward model, DPO directly optimizes the policy to increase the likelihood of preferred outputs and decrease the likelihood of rejected ones.\n",
    "\n",
    "## Step 1: Set up your environment\n",
    "\n",
    "### Dependencies\n",
    "First, ensure your environment has the correct libraries. Start with a pre-built container image and install LLaMA-Factory and DeepSpeed on top of it.\n",
    "\n",
    "Recommended container image:\n",
    "```bash\n",
    "anyscale/ray-llm:2.48.0-py311-cu128\n",
    "```\n",
    "\n",
    "Execute the following commands to install the required packages and optional tools for experiment tracking and faster model downloads:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Install the specific version of LLaMA-Factory\n",
    "pip install -q llamafactory==0.9.3\n",
    "\n",
    "# (Optional) For visualizing training metrics and logs\n",
    "pip install -q tensorboard==2.20.0\n",
    "\n",
    "# (Optional) For lightweight 8-bit and 4-bit optimizers and inference\n",
    "pip install -q bitsandbytes==0.47.0\n",
    "\n",
    "# (Optional) For AWQ quantization support\n",
    "pip install -q autoawq==0.2.9\n",
    "\n",
    "# (Optional) For accelerated model downloads from Hugging Face\n",
    "pip install -q hf_transfer==0.1.9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model and compute resources\n",
    "\n",
    "| Item | Value |\n",
    "|------|-------|\n",
    "| **Base model** | [`Qwen/Qwen2.5-7B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) |\n",
    "| **Workers** | 4 × L4 / A10G |\n",
    "\n",
    "Compared to SFT, DPO holds two copies of the model (policy and reference), and alignment datasets often use long contexts, so it's the ideal workflow for memory optimization techniques such as **QLoRA**. On 24 GB NVIDIA L4 GPUs, running DPO at FP16 for 7B models generally OOMs without QLoRA.\n",
    "\n",
    "## Step 2: Prepare the dataset\n",
    "\n",
    "### Understand the dataset\n",
    "This tutorial uses [`ultrafeedback.jsonl`](https://huggingface.co/datasets/kaitchup/UltraFeedback-prompt-chosen-rejected), a preference dataset tailored for DPO. Each sample contains one instruction **prompt** and two candidate completions: a **preferred** (`chosen`) response and a **less preferred** (`rejected`) response.\n",
    "\n",
    "This dataset includes:\n",
    "- `prompt`: An instruction or question to answer, often multi-sentence, with constraints.\n",
    "- `chosen`: The response that best follows the instruction.\n",
    "- `rejected`: A weaker alternative for the same prompt.\n",
    "\n",
    "**Dataset example**\n",
    "```json\n",
    "{\n",
    "  \"prompt\": \"Paraphrase the given questions to have different wording. Your paraphrased questions should have the same answer as the original question. Try to change the sentence as much as possible using synonyms and/or rearranging the structure of the sentence. The questions are in three domains: presidents, national parks, and dogs. Each question has a keyword indicating its domain. Keywords are \\\"this national park\\\", \\\"this dog breed\\\", and \\\"this president\\\", which will be replaced with the name of an actual president, a national park, or a breed of dog. Hence, in paraphrasing, this keyword should also be used the same way. Do not write questions that compare or involve multiple domains. Do not write open-ended or subjective questions (e.g., questions that can be answered differently by different people.) Make your questions specific and concrete. Your question should have the same type of answer as the original question(e.g., if the question is extractive, the paraphrased question should be extractive as well.)\\n\\nWhat lakes are in this national park?\",\n",
    "\n",
    "  \"rejected\": \"What bodies of water are located in this national park? \\n\\nWhich president is commonly known for his efforts to protect natural resources?\\n\\nWhich president is recognized for their dedication to preserving the environment? \\n\\nWhat type of dog breed is known for its loyalty and affectionate nature?\\n\\nWhat breed of dog is renowned for its faithfulness and loving personality?\",\n",
    "  \n",
    "  \"chosen\": \"Which bodies of water can be found within the borders of this particular national park?\"\n",
    "}\n",
    "```\n",
    "\n",
    "### Register the dataset\n",
    "\n",
    "To specify new datasets that are accessible across Ray worker nodes, you must first add a **`dataset_info.json`** to **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)** such as `/mnt/cluster_storage`. This configuration file acts as a central registry for all your datasets. It maps a custom name to your dataset file location, format, and column structure. \n",
    "\n",
    "If you plan to run DPO post-training on the `ultrafeedback` dataset, first complete the setup steps below. Ensure that you place the dataset files in a storage location that all workers can access like a shared mount or object storage. Avoid storing large files on the head node. \n",
    "\n",
    "`dataset_info.json`\n",
    "```json\n",
    "{\n",
    "  \"my_ultrafeedback\": {\n",
    "    \"file_name\": \"/mnt/cluster_storage/ultrafeedback.jsonl\",\n",
    "    \"ranking\": true,\n",
    "    \"columns\": {\n",
    "      \"prompt\": \"prompt\",\n",
    "      \"chosen\": \"chosen\",\n",
    "      \"rejected\": \"rejected\"\n",
    "    }\n",
    "  }\n",
    "}\n",
    "```\n",
    "\n",
    "For a more detailed dataset preparation and formatting guide, see [Choose your data format](https://docs.anyscale.com/llm/fine-tuning/data-preparation#preference-methods)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Make sure all files are accessible to worker nodes\n",
    "# Create a copy of the data in /mnt/cluster_storage\n",
    "wget https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/datasets/alpaca/ultrafeedback.jsonl -O /mnt/cluster_storage/ultrafeedback.jsonl\n",
    "# Create a copy of the dataset registry in /mnt/cluster_storage\n",
    "cp ../dataset-configs/dataset_info.json /mnt/cluster_storage/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create the preference-tuning config (DPO and QLoRA)\n",
    "\n",
    "Next, create the YAML configuration file that defines your DPO run. It specifies the base model, quantization (QLoRA), dataset, DPO hyperparameters, logging, and Ray cluster resources.\n",
    "\n",
    "**Important notes:**\n",
    "- **QLoRA quantization:** `quantization_bit: 4` with `quantization_method: bnb` applies quantization using bitsandbytes, reducing memory while preserving quality. If you use a model *pre-quantized* with AWQ, **omit** these keys.\n",
    "- **LoRA setup**: If you prefer standard LoRA, **disable quantization** by removing both `quantization_bit` and `quantization_method` from the config.\n",
    "- **Access & paths:** The YAML only needs to be on the **head node**, but any referenced paths (`dataset_dir`, `output_dir`) must reside on storage **reachable by all workers** (for example, `/mnt/cluster_storage/`).\n",
    "- **Gated models:** If your base model has gated access (for example, Llama) on Hugging Face, set `HF_TOKEN` in the runtime environment.\n",
    "\n",
    "### Configure LLaMA-Factory with Ray\n",
    "\n",
    "**Note**: To customize the training configuration, edit `train-configs/dpo_qlora.yaml`. \n",
    "\n",
    "```yaml\n",
    "# dpo_qlora.yaml\n",
    "\n",
    "### model\n",
    "trust_remote_code: true\n",
    "model_name_or_path: Qwen/Qwen2.5-7B-Instruct\n",
    "\n",
    "### method\n",
    "# If you instead want to use just LoRA, or a pre-quantized model like Qwen/Qwen2.5-7B-Instruct-AWQ, then omit the quantization_bit/method keys below\n",
    "quantization_bit: 4 # 4-bit base weights (QLoRA). Use 8 for 8-bit; omit for FP16/BF16\n",
    "quantization_method: bnb  # QLoRA via BitsAndBytes or hqq / eetq\n",
    "\n",
    "stage: dpo\n",
    "do_train: true\n",
    "finetuning_type: lora\n",
    "lora_rank: 8\n",
    "lora_target: all\n",
    "pref_beta: 0.1\n",
    "pref_loss: sigmoid  # choices: [sigmoid (dpo), orpo, simpo]\n",
    "\n",
    "# local dataset\n",
    "dataset: my_ultrafeedback\n",
    "dataset_dir: /mnt/cluster_storage\n",
    "\n",
    "template: qwen\n",
    "cutoff_len: 1024\n",
    "max_samples: 1000\n",
    "overwrite_cache: true\n",
    "preprocessing_num_workers: 16\n",
    "\n",
    "### output\n",
    "output_dir: qwen2.5_7b_qlora_dpo\n",
    "logging_steps: 5\n",
    "save_steps: 5              # For tensorboard logging purpose too. Can increase if not using tensorboard\n",
    "plot_loss: true\n",
    "report_to: tensorboard  # or none\n",
    "\n",
    "### train\n",
    "per_device_train_batch_size: 1\n",
    "gradient_accumulation_steps: 2\n",
    "num_train_epochs: 3.0  # Low for demo purpose; adjust as needed\n",
    "learning_rate: 5.0e-6\n",
    "bf16: true\n",
    "lr_scheduler_type: cosine\n",
    "warmup_ratio: 0.1\n",
    "ddp_timeout: 180000000\n",
    "\n",
    "### ray\n",
    "ray_run_name: qwen2.5_7b_qlora_dpo\n",
    "ray_storage_path: /mnt/cluster_storage/\n",
    "ray_num_workers: 4  # Number of GPUs to use.\n",
    "resources_per_worker:\n",
    "  GPU: 1\n",
    "  anyscale/accelerator_shape:4xL4: 0.001  # Use this to specify a specific node shape.\n",
    "  # accelerator_type:L4: 0.001            # Or use this to simply specify a GPU type.\n",
    "  # See https://docs.ray.io/en/master/ray-core/accelerator-types.html#accelerator-types for a full list of accelerator types.\n",
    "\n",
    "ray_init_kwargs:\n",
    "  runtime_env:\n",
    "    env_vars:\n",
    "      # If using gated models like meta-llama/Llama-3.1-8B-Instruct\n",
    "      # HF_TOKEN: <your_huggingface_token>\n",
    "      # Enable faster downloads if hf_transfer is installed:\n",
    "      HF_HUB_ENABLE_HF_TRANSFER: '1'\n",
    "```\n",
    "\n",
    "## Step 4: Train and monitor\n",
    "\n",
    "With all configurations in place, you can launch fine-tuning or post-training in one of two ways:\n",
    "\n",
    "### Option A: Run from a workspace (quick start)\n",
    "\n",
    "The `USE_RAY=1` prefix tells LLaMA-Factory to run in distributed mode on the Ray cluster attached to your workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "USE_RAY=1 llamafactory-cli train ../train-configs/dpo_qlora.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option B: Run as an Anyscale job (production)\n",
    "\n",
    "For longer or production runs, submit the training as an **Anyscale job**. Jobs run outside your interactive session for better stability, retries, and durable logs. Package LLaMA-Factory and other libraries in a container image and launch with a short job config. See [Run LLaMA-Factory as an Anyscale job](https://docs.anyscale.com/llm/fine-tuning/llamafactory-jobs) for the step-by-step guide.\n",
    "\n",
    "### Tracking with TensorBoard\n",
    "If you enabled TensorBoard logging (`report_to: tensorboard` in your YAML), you can watch metrics (for example, training loss) update live and compare multiple runs with the same run name side-by-side.\n",
    "\n",
    "- **While the job is running:** LLaMA-Factory prints a ready-to-run command that starts with `tensorboard --logdir`. Open a new terminal and run it. For example:\n",
    "  ```bash\n",
    "  tensorboard --logdir /tmp/ray/session_*/artifacts/*/qwen2.5_7b_qlora_dpo/driver_artifacts\n",
    "  ```\n",
    "\n",
    "- **After the job:** Point TensorBoard at `{ray_storage_path}/{ray_run_name}/`. Each `TorchTrainer_*` subfolder holds event files for a single run. Using the parent folder aggregates all runs for easy comparison.\n",
    "  ```bash\n",
    "  tensorboard --logdir /mnt/cluster_storage/qwen2.5_7b_qlora_dpo\n",
    "  ```\n",
    "\n",
    "In your Anyscale workspace, look for the open **port 6006** labeled **TensorBoard** to view the dashboards.\n",
    "\n",
    "![Anyscale workspace showing open ports with TensorBoard on port 6006](https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/open-ports.png)\n",
    "\n",
    "**TensorBoard example**\n",
    "\n",
    "![TensorBoard](https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/3.2.2/3.2.2-tensorboard.png)\n",
    "\n",
    "For a more detailed guide on tracking experiments with other tools such as Weights & Biases or MLflow, see [Observability and tracking](https://docs.anyscale.com/llm/fine-tuning/observability-and-tracking).\n",
    "\n",
    "## Step 5: Locate checkpoints\n",
    "\n",
    "Ray Train writes checkpoints under `ray_storage_path/ray_run_name`. In this example run, the path is: `/mnt/cluster_storage/qwen2.5_7b_qlora_dpo`. \n",
    "\n",
    "Inside, you see a **trainer session** directory named like:\n",
    "`TorchTrainer_ff224_00000_0_2025-09-19_15-57-20/`.\n",
    "\n",
    "- Ray Train creates `TorchTrainer_*` **when the trainer starts**; the suffix encodes a short run ID and the **start timestamp**.\n",
    "- Within that directory, Ray Train names checkpoints `checkpoint_000xxx/`, where the number is the saved ordered checkpoints.\n",
    "\n",
    "Control the save cadence with `save_strategy` and `save_steps`. For instructions on how to resume interrupted training with `resume_from_checkpoint` and more, see [Understand the artifacts directory](https://docs.anyscale.com/llm/fine-tuning/checkpointing#artifacts-directory).\n",
    "\n",
    "## Step 6: Export the model\n",
    "\n",
    "If you use LoRA, you can keep the base model and adapters separate for [multi-LoRA deployment](https://docs.anyscale.com/llm/serving/multi-lora) or [merge the adapters](https://docs.anyscale.com/llm/fine-tuning/checkpointing#merge-lora) into the base model for low-latency inference. \n",
    "\n",
    "For full fine-tuning or freeze-tuning, export the fine-tuned model directly.\n",
    "\n",
    "You may optionally apply [post-training quantization](https://docs.anyscale.com/llm/fine-tuning/checkpointing#ptq) on merged or full models before serving."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
